package com.jay.net.clazz.socket.udp;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.io.InterruptedIOException;
import java.net.DatagramPacket;
import java.net.DatagramSocket;
import java.net.InetAddress;
import java.net.SocketAddress;
import java.net.SocketException;

//All times are expressed in seconds.
//ReliabilityConstants interface, just defines constants.
interface ReliabilityConstants {
	// Timeout minima/maxima
	public static final int MIN_RETRANSMIT_TIMEOUT = 1;
	public static final int MAX_RETRANSMIT_TIMEOUT = 64;
	// Maximum retransmissions per datagram, suggest 3 or 4.
	public static final int MAX_RETRANSMISSIONS = 4;
}

// RoundTripTimer class.
class RoundTripTimer implements ReliabilityConstants {
	float roundTripTime = 0.0f;// most recent RTT
	float smoothedTripTime = 0.0f;// smoothed RTT
	float deviation = 0.75f; // smoothed mean deviation
	short retransmissions = 0;// retransmit count: 0, 1, 2, …
	// current retransmit timeout
	float currentTimeout = minmax(calculateRetransmitTimeout());

	/** @return the re-transmission timeout. */
	private int calculateRetransmitTimeout() {
		return (int) (smoothedTripTime + 4.0 * deviation);
	}

	/** @return the bounded retransmission timeout. */
	private float minmax(float rto) {
		return Math.min(Math.max(rto, MIN_RETRANSMIT_TIMEOUT),
				MAX_RETRANSMIT_TIMEOUT);
	}

	/** Called before each new packet is transmitted. */
	void newPacket() {
		retransmissions = 0;
	}

	/**
	 * @return the timeout for the packet.
	 */
	float currentTimeout() {
		return currentTimeout;
	}

	/**
	 * Called straight after a successful receive. Calculates the round-trip
	 * time, then updates the smoothed round-trip time and the variance
	 * (deviation).
	 * 
	 * @param ms
	 *            time in ms since starting the transmission.
	 */
	void stoppedAt(long ms) {
		// Calculate the round-trip time for this packet.
		roundTripTime = ms / 1000;
		// Update our estimators of round-trip time
		// and its mean deviation.
		double delta = roundTripTime - smoothedTripTime;
		smoothedTripTime += delta / 8.0;
		deviation += (Math.abs(delta) - deviation) / 4.0;
		// Recalculate the current timeout.
		currentTimeout = minmax(calculateRetransmitTimeout());
	}

	/**
	 * Called after a timeout has occurred.
	 * 
	 * @return true if it's time to give up, false if we can retransmit.
	 */
	boolean isTimeout() {
		currentTimeout *= 2; // next retransmit timeout
		retransmissions++;
		return retransmissions > MAX_RETRANSMISSIONS;
	}
} // RoundTripTimer class

// ReliableDatagramSocket class
public class ReliableDatagramSocket extends DatagramSocket implements
		ReliabilityConstants {
	RoundTripTimer roundTripTimer = new RoundTripTimer();
	private boolean reinit = false;
	private long sendSequenceNo = 0; // send sequence #
	private long recvSequenceNo = 0; // recv sequence #
	/* anonymous initialization for all constructors */
	{
		init();
	}

	/**
	 * Construct a ReliableDatagramSocket
	 * 
	 * @param port
	 *            Local port: reeive on any interface/address
	 * @exception SocketException
	 *                can't create the socket
	 */
	public ReliableDatagramSocket(int port) throws SocketException {
		super(port);
	}

	/**
	 * Construct a ReliableDatagramSocket
	 * 
	 * @param port
	 *            Local port
	 * @param localAddr
	 *            local interface address to use
	 * @exception SocketException
	 *                can't create the socket
	 */
	public ReliableDatagramSocket(int port, InetAddress localAddr)
			throws SocketException {
		super(port, localAddr);
	}

	/**
	 * Construct a ReliableDatagramSocket, JDK >= 1.4.
	 * 
	 * @param localAddr
	 *            local socket address to use
	 * @exception SocketException
	 *                can't create the socket
	 */
	public ReliableDatagramSocket(SocketAddress localAddr)
			throws SocketException {
		super(localAddr);
	}

	/**
	 * Overrides DatagramSocket.connect(): Does the connect, then
	 * (re-)initializes the statistics for the connection.
	 * 
	 * @param dest
	 *            Destination address
	 * @param port
	 *            Destination port
	 */
	public void connect(InetAddress dest, int port) {
		super.connect(dest, port);
		init();
	}

	/**
	 * Overrides JDK 1.4 DatagramSocket.connect(). Does the connect, then
	 * (re-)initializes the statistics for the connection.
	 * 
	 * @param dest
	 *            Destination address
	 */
	public void connect(SocketAddress dest) throws SocketException {
		super.connect(dest);
		init();
	}

	/** Initialize */
	private void init() {
		this.roundTripTimer = new RoundTripTimer();
	}

	/**
	 * Send and receive reliably, retrying adaptively with exponential backoff
	 * until the response is received or timeout occurs.
	 * 
	 * @param sendPacket
	 *            outgoing request datagram
	 * @param recvPacket
	 *            incoming reply datagram
	 * @exception IOException
	 *                on any error
	 * @exception InterruptedIOException
	 *                on timeout
	 */
	public synchronized void sendReceive(DatagramPacket sendPacket,
			DatagramPacket recvPacket) throws IOException,
			InterruptedIOException {
		// re-initialize after timeout
		if (reinit) {
			init();
			reinit = false;
		}
		roundTripTimer.newPacket();
		long start = System.currentTimeMillis();
		long sequenceNumber = getSendSequenceNo();
		// Loop until final timeout or some unexpected exception
		for (;;) {
			// keep using the same sequenceNumber while retrying
			setSendSequenceNo(sequenceNumber);
			send(sendPacket);// may throw
			int timeout = (int) (roundTripTimer.currentTimeout() * 1000.0 + 0.5);
			long soTimeoutStart = System.currentTimeMillis();
			try {
				for (;;) {
					// Adjust socket timeout for time already elapsed
					int soTimeout = timeout
							- (int) (System.currentTimeMillis() - soTimeoutStart);
					setSoTimeout(soTimeout);
					receive(recvPacket);
					long recvSequenceNumber = getRecvSequenceNo();
					if (recvSequenceNumber == sequenceNumber) {
						// Got the correct reply:
						// stop timer, calculate new RTT values
						long ms = System.currentTimeMillis() - start;
						roundTripTimer.stoppedAt(ms);
						return;
					}
				}
			} catch (InterruptedIOException exc) {
				// timeout: retry?
				if (roundTripTimer.isTimeout()) {
					reinit = true;
					// rethrow InterruptedIOException to caller
					throw exc;
				}
				// else continue
			}
			// may throw other SocketException or IOException
		} // end re-transmit loop
	} // sendReceive()

	/**
	 * @return the last received sequence number; used by servers to obtain the
	 *         reply sequenceNumber.
	 */
	public long getRecvSequenceNo() {
		return recvSequenceNo;
	}

	/** @return the last sent sequence number */
	private long getSendSequenceNo() {
		return sendSequenceNo;
	}

	/**
	 * Set the next send sequence number. Used by servers to set the reply
	 * sequenceNumber from the received packet:
	 * 
	 * . * socket.setSendSequenceNo(socket.getRecvSequenceNo());
	 * 
	 * @param sendSequenceNo
	 *            Next sequence number to send.
	 */
	public void setSendSequenceNo(long sendSequenceNo) {
		this.sendSequenceNo = sendSequenceNo;
	}

	/**
	 * override for DatagramSocket.receive: handles the sequence number.
	 * 
	 * @param packet
	 *            DatagramPacket
	 * @exception IOException
	 *                I/O error
	 */
	public void receive(DatagramPacket packet) throws IOException {
		super.receive(packet);
		// read sequence number and remove it from the packet
		ByteArrayInputStream bais = new ByteArrayInputStream(packet.getData(),
				packet.getOffset(), packet.getLength());
		DataInputStream dis = new DataInputStream(bais);
		recvSequenceNo = dis.readLong();
		byte[] buffer = new byte[dis.available()];
		dis.read(buffer);
		packet.setData(buffer, 0, buffer.length);
	}

	/**
	 * override for DatagramSocket.send: handles the sequence number.
	 * 
	 * @param packet
	 *            DatagramPacket
	 * @exception IOException
	 *                I/O error
	 */
	public void send(DatagramPacket packet) throws IOException {
		ByteArrayOutputStream baos = new ByteArrayOutputStream();
		DataOutputStream dos = new DataOutputStream(baos);
		// Write the sequence number, then the user data.
		dos.writeLong(sendSequenceNo++);
		dos.write(packet.getData(), packet.getOffset(), packet.getLength());
		dos.flush();
		// Construct a new packet with this new data and send it.
		byte[] data = baos.toByteArray();
		packet = new DatagramPacket(data, baos.size(), packet.getAddress(),
				packet.getPort());
		super.send(packet);
	}
} // end of ReliableDatagramSocket class
